import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Function;

/**
 * 请求合并处理器
 * 在高并发场景下，将短时间内的相同请求合并处理，减轻系统负担
 * @param <K> 请求的键类型
 * @param <V> 请求的返回值类型
 */
public class RequestMerger<K, V> {
    // 存储正在处理中的请求
    private final Map<K, RequestWrapper<V>> ongoingRequests = new ConcurrentHashMap<>();
    private final Queue<RequestWrapper<V>> requestQueue = new LinkedBlockingQueue<>();
    private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(2);
    private final long timeoutMillis;
    private final int maxBatchSize;
    
    // 包装请求，添加时间戳和分组信息
    private static class RequestWrapper<V> {
        final CompletableFuture<V> future;
        final long timestamp;
        final Object params;
        final String groupKey;  // 添加分组键
        
        RequestWrapper(Object params, String groupKey) {
            this.future = new CompletableFuture<>();
            this.timestamp = System.currentTimeMillis();
            this.params = params;
            this.groupKey = groupKey;
        }
        
        boolean isExpired(long timeoutMillis) {
            return System.currentTimeMillis() - timestamp > timeoutMillis;
        }
    }
    
    public RequestMerger(long timeoutMillis, int maxBatchSize) {
        this.timeoutMillis = timeoutMillis;
        this.maxBatchSize = maxBatchSize;
        
        // 定期批量处理请求
        scheduler.scheduleAtFixedRate(this::processBatchRequests, 
                100, 10, TimeUnit.MILLISECONDS);
                
        // 定期清理过期请求
        scheduler.scheduleAtFixedRate(this::cleanExpiredRequests,
                timeoutMillis / 2, timeoutMillis / 2, TimeUnit.MILLISECONDS);
    }
    
    public CompletableFuture<V> execute(K key, Object params, String groupKey, Function<List<Object>, Map<K, V>> batchFunction) {
        // 创建新请求，包含分组信息
        RequestWrapper<V> wrapper = new RequestWrapper<>(params, groupKey);
        
        // 检查是否有未过期的相同请求
        RequestWrapper<V> existing = ongoingRequests.get(key);
        if (existing != null && !existing.isExpired(timeoutMillis) && !existing.future.isDone()) {
            return existing.future;
        }
        
        // 放入请求映射和队列
        ongoingRequests.put(key, wrapper);
        requestQueue.offer(wrapper);
        
        return wrapper.future;
    }
    
    private void processBatchRequests() {
        int size = requestQueue.size();
        if (size == 0) {
            return;
        }
        
        // 按分组收集请求
        Map<String, List<RequestWrapper<V>>> groupedRequests = new HashMap<>();
        Map<String, List<Object>> groupedParams = new HashMap<>();
        Map<String, Map<Object, RequestWrapper<V>>> groupedRequestMaps = new HashMap<>();
        
        // 收集并分组请求
        for (int i = 0; i < Math.min(size, maxBatchSize); i++) {
            RequestWrapper<V> request = requestQueue.poll();
            if (request != null && !request.isExpired(timeoutMillis)) {
                // 按groupKey分组
                groupedRequests.computeIfAbsent(request.groupKey, k -> new ArrayList<>()).add(request);
                groupedParams.computeIfAbsent(request.groupKey, k -> new ArrayList<>()).add(request.params);
                groupedRequestMaps.computeIfAbsent(request.groupKey, k -> new HashMap<>())
                    .put(request.params, request);
            }
        }
        
        // 按分组处理请求
        groupedRequests.forEach((groupKey, requests) -> {
            if (requests.isEmpty()) {
                return;
            }
            
            try {
                // 批量处理同一分组的请求
                Map<K, V> results = batchFunction.apply(groupedParams.get(groupKey));
                
                // 分发结果
                Map<Object, RequestWrapper<V>> requestMap = groupedRequestMaps.get(groupKey);
                results.forEach((key, value) -> {
                    RequestWrapper<V> request = requestMap.get(key);
                    if (request != null) {
                        request.future.complete(value);
                    }
                });
            } catch (Exception e) {
                // 处理异常
                requests.forEach(request -> request.future.completeExceptionally(e));
            }
        });
    }
    
    private void cleanExpiredRequests() {
        // 清理过期请求
        ongoingRequests.entrySet().removeIf(entry -> {
            RequestWrapper<V> wrapper = entry.getValue();
            if (wrapper.isExpired(timeoutMillis) && !wrapper.future.isDone()) {
                wrapper.future.completeExceptionally(
                    new TimeoutException("Request expired after " + timeoutMillis + "ms"));
                return true;
            }
            return wrapper.future.isDone();
        });
    }
    
    public void shutdown() {
        scheduler.shutdown();
    }
}